﻿using System;
using System.Collections.Generic;
using System.Linq;
using UnityEditor;
using UnityEngine;

// 树形结构数据项（增加层级和父节点引用）
[Serializable]
public class TreeItem
{
    public string name;
    public object data;
    public List<TreeItem> children = new List<TreeItem>();
    public bool isExpanded = false;
    public TreeItem parent;  // 父节点引用
    public int level;       // 节点层级（0=根，1=Shader，2=材质，3=对象）

    public TreeItem(string name, object data = null, TreeItem parent = null, int level = 0)
    {
        this.name = name;
        this.data = data;
        this.parent = parent;
        this.level = level;
    }
}

// 材质使用信息
[Serializable]
public class MaterialUsageInfo
{
    public Material material;
    public int usageCount;
    public List<GameObject> usingObjects = new List<GameObject>();
}

public class ShaderMaterialAnalyzer : EditorWindow
{
    private TreeItem rootItem;
    private Vector2 scrollPosition;
    private Dictionary<Shader, List<Material>> shaderToMaterials = new Dictionary<Shader, List<Material>>();
    private Dictionary<Material, MaterialUsageInfo> materialUsage = new Dictionary<Material, MaterialUsageInfo>();
    private bool isScanning = false;

    [MenuItem("TA/辅助/Shader & Material Analyzer")]
    public static void ShowWindow()
    {
        GetWindow<ShaderMaterialAnalyzer>("Shader与材质分析器");
    }

    private void OnGUI()
    {
        GUILayout.Label("场景Shader与材质统计", EditorStyles.boldLabel);

        if (isScanning)
        {
            EditorGUILayout.LabelField("正在扫描场景...", EditorStyles.centeredGreyMiniLabel);
            return;
        }

        if (GUILayout.Button("扫描当前场景"))
        {
            ScanScene();
        }

        if (rootItem != null)
        {
            int totalShaders = shaderToMaterials.Count;
            int totalMaterials = materialUsage.Count;
            int totalObjects = GameObject.FindObjectsOfType<Renderer>().Length;

            EditorGUILayout.LabelField($"统计结果: {totalShaders}个Shader, {totalMaterials}个材质, {totalObjects}个渲染对象");

            scrollPosition = EditorGUILayout.BeginScrollView(scrollPosition);
            DrawTreeView(rootItem);
            EditorGUILayout.EndScrollView();
        }
        else
        {
            EditorGUILayout.HelpBox("点击上方按钮扫描场景以获取数据", MessageType.Info);
        }
    }

    private void ScanScene()
    {
        isScanning = true;
        EditorApplication.RepaintHierarchyWindow();

        shaderToMaterials.Clear();
        materialUsage.Clear();

        Renderer[] renderers = GameObject.FindObjectsOfType<Renderer>();
        int batchSize = 500;

        for (int i = 0; i < renderers.Length; i += batchSize)
        {
            int end = Mathf.Min(i + batchSize, renderers.Length);
            for (int j = i; j < end; j++)
            {
                ProcessRenderer(renderers[j]);
            }

            EditorApplication.QueuePlayerLoopUpdate();
        }

        BuildTreeStructure();
        isScanning = false;
    }

    private void ProcessRenderer(Renderer renderer)
    {
        Material[] materials = renderer.sharedMaterials;

        for (int i = 0; i < materials.Length; i++)
        {
            Material material = materials[i];
            if (material == null) return;

            if (!materialUsage.ContainsKey(material))
            {
                materialUsage[material] = new MaterialUsageInfo
                {
                    material = material,
                    usageCount = 0,
                    usingObjects = new List<GameObject>()
                };
            }

            MaterialUsageInfo info = materialUsage[material];
            info.usageCount++;

            if (!info.usingObjects.Contains(renderer.gameObject))
            {
                info.usingObjects.Add(renderer.gameObject);
            }

            if (!shaderToMaterials.ContainsKey(material.shader))
            {
                shaderToMaterials[material.shader] = new List<Material>();
            }

            if (!shaderToMaterials[material.shader].Contains(material))
            {
                shaderToMaterials[material.shader].Add(material);
            }
        }
    }

    private void BuildTreeStructure()
    {
        rootItem = new TreeItem("场景资源统计", level: 0);

        foreach (var pair in shaderToMaterials.OrderBy(x => x.Key.name))
        {
            Shader shader = pair.Key;
            List<Material> materials = pair.Value;

            // ✅ 为Shader节点赋值data字段（存储Shader对象）
            TreeItem shaderItem = new TreeItem($"{shader.name} ({materials.Count}个材质)",
                                             shader, parent: rootItem, level: 1);
            rootItem.children.Add(shaderItem);

            foreach (Material material in materials.OrderBy(x => x.name))
            {
                if (materialUsage.TryGetValue(material, out MaterialUsageInfo info))
                {
                    TreeItem materialItem = new TreeItem($"{material.name} ({info.usageCount}个对象使用)",
                                                     material, shaderItem, 2);
                    shaderItem.children.Add(materialItem);

                    // 延迟加载对象节点
                    materialItem.children.Add(new TreeItem("点击展开查看使用对象", info, materialItem, 3));
                }
            }
        }
    }

    private void DrawTreeView(TreeItem item)
    {
        int indentLevel = item.level;
        using (new EditorGUI.IndentLevelScope(indentLevel))
        {
            bool wasExpanded = item.isExpanded;

            if (item.level >= 1)  // 仅对二级及以下节点启用互斥
            {
                item.isExpanded = EditorGUILayout.Foldout(wasExpanded, item.name);

                if (item.isExpanded && wasExpanded == false)
                {
                    CloseSiblingNodes(item);
                }
            }
            else
            {
                item.isExpanded = EditorGUILayout.Foldout(wasExpanded, item.name);
            }

            Rect clickRect = GUILayoutUtility.GetLastRect();
            if (clickRect.Contains(Event.current.mousePosition) &&
                Event.current.type == EventType.MouseDown &&
                Event.current.button == 0)
            {
                HandleItemClick(item);
                Event.current.Use();
            }

            if (item.isExpanded && item.children.Count > 0)
            {
                if (item.level == 2)  // 材质节点（二级）展开时加载对象
                {
                    LoadObjectChildren(item);
                }

                foreach (TreeItem child in item.children)
                {
                    DrawTreeView(child);
                }
            }
        }
    }

    private void CloseSiblingNodes(TreeItem currentItem)
    {
        if (currentItem.parent == null) return;

        foreach (TreeItem sibling in currentItem.parent.children)
        {
            if (sibling != currentItem && sibling.isExpanded)
            {
                sibling.isExpanded = false;
            }
        }
    }

    private void LoadObjectChildren(TreeItem materialItem)
    {
        if (materialItem.children.Count == 1 && materialItem.children[0].data is MaterialUsageInfo)
        {
            MaterialUsageInfo info = (MaterialUsageInfo)materialItem.children[0].data;
            materialItem.children.Clear();

            foreach (GameObject obj in info.usingObjects.OrderBy(x => x.name))
            {
                TreeItem objectItem = new TreeItem(obj.name, obj, materialItem, 3);
                materialItem.children.Add(objectItem);
            }
        }
    }

    private void HandleItemClick(TreeItem item)
    {
        if (item.data == null) return;

        // ✅ 新增Shader定位逻辑
        if (item.data is Shader shader)
        {
            EditorGUIUtility.PingObject(shader);
            Selection.activeObject = shader;
            EditorApplication.ExecuteMenuItem("Window/Project");
        }
        else if (item.data is Material material)
        {
            EditorGUIUtility.PingObject(material);
            Selection.activeObject = material;
            EditorApplication.ExecuteMenuItem("Window/Project");
        }
        else if (item.data is GameObject gameObject)
        {
            EditorGUIUtility.PingObject(gameObject);
            Selection.activeObject = gameObject;
            SceneView.lastActiveSceneView.FrameSelected(gameObject);
            EditorApplication.ExecuteMenuItem("Window/Scene");
        }
        else if (item.data is MaterialUsageInfo && item.level == 3)
        {
            item.isExpanded = !item.isExpanded;
            if (item.isExpanded) CloseSiblingNodes(item);
        }
    }
}